# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/29 11:24 
@Author  : ZhangShenao 
@File    : keyword_text_splitter.py 
@Desc    : 自定义关键词文档分割器

实现自定义文档分割器,需要继承TextSplitter类,并实现split_text()方法

"""
from typing import List

import jieba.analyse
from langchain_text_splitters import TextSplitter


class KeywordTextSplitter(TextSplitter):
    """
    关键词文档分割器
    按照指定的分割符分割文档,并且对每个分片提取关键词
    最终返回每个分片的关键词列表
    """

    def __init__(self, separator: str, top_k: int):
        """
        构造函数
        :param separator: 文档分隔符
        :param top_k: 每个分片提取的关键词数量
        """

        # 调用父类的构造函数
        super().__init__()

        # 初始化
        self.__separator = separator
        self.__top_k = top_k

    def split_text(self, text: str) -> List[str]:
        # 按照分隔符分割文档
        chunks = text.split(self.__separator)

        # 针对每个文档分片,提取top_k个关键词
        key_words = []
        for chunk in chunks:
            tags = jieba.analyse.extract_tags(chunk, topK=self.__top_k)
            key_words.append(tags)

        # 返回逐个分片的关键词列表
        return [','.join(keyword) for keyword in key_words]
